args=commandArgs(T)

XCRDirListi=args[1]
outXCRSeqFile=args[2]
XCRflag=args[3]  #1,TCR;2,BCR

#finish in 2021.11.11
#output:table("key","conid","fa").TCR key: barcode:V:D:J:CDR3   BCR key:barcode:V:D:J:CDR3

#XCRDirListi="/home/xingjun/zdata1/SCWork/P4/P4.TCR.list"
#outXCRSeqFile="/home/xingjun/zdata1/aaa.xls"
#XCRflag="1"

#base function
readList <- function(LTable){
  a <- read.table(LTable,sep = "\t")
  sampleList  <- as.vector(a$V1)
  b <- as.vector(a$V2)
  names(b) <- as.vector(a$V1)
  return(b)
}

getConSeq <- function(conseq.loc)
{
  #read consensus.fasta
  conseq <- readLines(conseq.loc)
  if(length(conseq) %% 2 == 0)
  {
    idVec <- c()
    seqVec <- c()
    
    
    for(i in seq(1,length(conseq),by=2))
    {
      id <- conseq[i]
      seq <- conseq[i+1]
      idVec <- c(idVec,id)
      seqVec <- c(seqVec,seq)
      
    }
    
    idVec <- gsub(">","",idVec)
    conidSeqTable <- cbind(idVec,seqVec)
    colnames(conidSeqTable) <- c("conid","seq")
    return(conidSeqTable)
    
  }else
  {
    print("num error!")
    return(FALSE)
  }
  
}

OneGetKeySeq <- function(celltype,contig.loc,conseq.loc,flag)
{
#need function:getConSeq(return a table(conid,seq))
#input:flag:1,2;celltype:Tumor,Normal,...

  #read barcode contig table
  contig <- read.csv(contig.loc)
  
  #rename barcode
  all.cell <- as.vector(contig[,1])
  all.cell <- paste(celltype,all.cell,sep = "")
  contig[,1] <- all.cell
  #rename barcode end.
  
  #read barcode contig table end.
  
  #read conid-seq file
  conSeq <- getConSeq(conseq.loc)  #conid - fa
  conSeq <- as.data.frame(conSeq)
  rownames(conSeq) <- conSeq$conid
  #read conid-seq file end.
  
  #cobime conid and seq
  if(flag == 1) #TCR
  {
    contig1 <- contig
    contig1 <- contig1[contig1$high_confidence=="True",]
    contig1 <- contig1[contig1$chain=="TRA" | contig1$chain=="TRB",]
    
    conid <- as.vector(contig1$raw_consensus_id)
    key <- paste(contig1$barcode,contig1$v_gene,contig1$d_gene,contig1$j_gene,contig1$cdr3,sep=":")
    keyConTable <- cbind(key,conid)
    keyConTable <- as.data.frame(keyConTable)
    
  }else if(flag == 2) #BCR
  {
    contig1 <- contig
    contig1 <- contig1[contig1$high_confidence == "True",]
    contig1 <- contig1[contig1$chain == "IGH" | contig1$chain == "IGK" | contig1$chain == "IGL",]
    
    conid <- as.vector(contig1$raw_consensus_id)
    key <- paste(contig1$barcode,contig1$v_gene,contig1$d_gene,contig1$j_gene,contig1$cdr3,sep=":")
    keyConTable <- as.data.frame(cbind(key,conid))
  }else
  {
    print("flag error!!!")
    return(FALSE)
  }
  
  #add seq col
  seqLines <- c()
  for(i in keyConTable$conid)
  {
    seqLines <- c(seqLines,as.character(conSeq[i,]$seq))
  }
  keyConTable <- cbind(keyConTable,seqLines)
  #add seq col end.
  
  return(keyConTable)
  #cobime conid and seq end.
  
}
#base function end.

#main
a <- readList(XCRDirListi)
cellTypeList <- names(a)

FinalTable <- NULL
for(i in seq(1,length(cellTypeList),1))
{
  celltype <- cellTypeList[i]
  dataDir <- as.character(a[i])
  print(celltype)
  
  contig.loc <- paste0(dataDir,"/filtered_contig_annotations.csv")
  conseq.loc <- paste0(dataDir,"/consensus.fasta")
  
  FinalTable <- rbind(FinalTable,OneGetKeySeq(celltype,contig.loc,conseq.loc,XCRflag))
  
}

colnames(FinalTable) <- c("key","conid","fa")
write.table(FinalTable,file=outXCRSeqFile,row.names = F,col.names = T,sep="\t",quote=F)
#main end.
